# pathformer.py
import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import Dataset, DataLoader


# 1. Patch Embedding（片段嵌入）
class PatchEmbedding(nn.Module):
    """增强稳定性的片段嵌入模块"""
    def __init__(self, input_dim, patch_size, embed_dim):
        super().__init__()
        self.patch_size = patch_size
        self.projection = nn.Linear(input_dim * patch_size, embed_dim)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        batch_size, seq_len, feat_dim = x.size()
        # 动态调整序列长度确保可被patch_size整除
        num_patches = seq_len // self.patch_size
        valid_len = num_patches * self.patch_size
        x = x[:, :valid_len, :] # 自动截断
        x = x.view(batch_size, num_patches, feat_dim * self.patch_size)
        x = self.projection(x)
        return self.norm(x)


# 2. Positional Encoding（位置编码）
class PositionalEncoding(nn.Module):
    """增强型位置编码"""
    def __init__(self, embed_dim, max_len=5000):
        super().__init__()
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-torch.log(torch.tensor(10000.0)) / embed_dim))
        pe = torch.zeros(1, max_len, embed_dim)
        pe[0, :, 0::2] = torch.sin(position * div_term)
        pe[0, :, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1), :]


# 3. pathformer 模型（适配多特征预测）
class PathFormer(pl.LightningModule):
    """整合残差连接的时空预测模型"""

    def __init__(self, input_dim, patch_size=30, embed_dim=512,
                 num_heads=8, num_layers=6, output_dim=15,
                 dropout=0.2, ff_dim=2048):
        super().__init__()
        self.save_hyperparameters()

        # 残差连接优化
        if input_dim != output_dim:
            self.residual = nn.Sequential(
                nn.Linear(input_dim, embed_dim),
                nn.Linear(embed_dim, output_dim))
        else:
            self.residual = None

            # 核心模块
            self.patch_embed = PatchEmbedding(input_dim, patch_size, embed_dim)
            self.pos_encoder = PositionalEncoding(embed_dim)
            self.transformer = nn.TransformerEncoder(
                encoder_layer=nn.TransformerEncoderLayer(
                    d_model=embed_dim,
                    nhead=num_heads,
                    dim_feedforward=ff_dim,
                    dropout=dropout,
                    batch_first=True,
                    activation='gelu'
                ),
                num_layers=num_layers
            )
            self.fc = nn.Linear(embed_dim, output_dim)
            self.dropout = nn.Dropout(dropout)
            self._init_weights()

    def _init_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_normal_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

    def forward(self, x):
        # 残差连接优化
        residual = x[:, -1, :] if self.residual else None  # 取最后一个时间步

        x = self.patch_embed(x)
        x = self.pos_encoder(x)
        x = self.transformer(x)
        x = x.mean(dim=1)
        x = self.dropout(x)
        output = self.fc(x)

        if self.residual:
            output += self.residual(residual)
        return output

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.HuberLoss(delta=1.0)(y_hat, y)
        self.log('train_loss', loss, prog_bar=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(),
            lr=1e-4,
            weight_decay=0.01
        )
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            patience=10,
            mode='min'
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "train_loss",
            }
        }


# 4. 数据集类（适配多目标）
class TimeSeriesDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.as_tensor(X, dtype=torch.float32)
        self.y = torch.as_tensor(y, dtype=torch.float32)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


# 5. 数据加载器（简化版）
# def create_dataloaders(X_train, y_train, batch_size=32, workers=4):
#     """与run.py训练流程兼容的数据加载器"""
#     dataset = TimeSeriesDataset(X_train, y_train)
#     return DataLoader(
#         dataset,
#         batch_size=batch_size,
#         shuffle=True,
#         num_workers=workers,
#         persistent_workers=workers > 0,
#         pin_memory=True  # 加速GPU传输
#     )